from langchain_community.utilities import SQLDatabase
from langchain_experimental.sql import SQLDatabaseChain
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
db_name = 'offer_db.sqlite'
db = SQLDatabase.from_uri(f"sqlite:///{db_name}")
# 创建OpenAI LLM客户端
from config.model_config import get_chat_openai_zhipu_flash_250414
llm = get_chat_openai_zhipu_flash_250414(temperature=0)

db_chain = SQLDatabaseChain.from_llm(llm, db)

PROMPT_TEMPLATE = """
                你会接收到一个查询，你的任务是从`offer_retailer`表中的`OFFER`字段检索相关offer。
                查询可能是混合大小写的，所以也要搜索大写版本的查询。
                重要的是，你可能需要使用数据库中其他表的信息，即：`brand_category`, `categories`, `offer_retailer`，来检索正确的offer。
                不要虚构offer。如果在`offer_retailer`表中找不到offer，返回字符串：`NONE`。
                如果你能从`offer_retailer`表中检索到offer，用分隔符`#`分隔每个offer。例如，输出应该是这样的：`offer1#offer2#offer3`。
                注意生成的sql必须不要带上```sql```。
                如果SQLResult为空，返回`None`。不要生成任何offer。
                这是查询：`{}`
                """
query = "Spend"
prompt = PROMPT_TEMPLATE.format(query)

# 运行SQL数据库链以检索offer
retrieved_offers = db_chain.run(prompt)
print(retrieved_offers)

retrieved_offers_list = retrieved_offers.split('#')
for offer in retrieved_offers_list:
    print(offer)
